import numpy as np



class HMM(object):
    def __init__(self, N, M, pi=None, A=None, B=None):
        self.N = N # 对应状态数量，盒子的数量
        self.M = M # 对应的是观测数量，球颜色的数量
        self.pi = pi # 初始概率向量
        self.A = A # 转移概率矩阵
        self.B = B # 观测概率矩阵

    
    def get_data_with_distribute(self, dist):
        # 根据给定的概率分布随机返回数据（索引）
        return np.random.choice(np.arange(len(dist)), p=dist)
    

    def generate(self, T:int):
        '''
        根据给定的参数生成观测序列
        T: 指定要生成数据的数量
        '''
        # 根据初始概率分布生成第一个状态，即从哪个盒子取球
        z = self.get_data_with_distribute(self.pi)
        # 生成第一个观测数据，即取出的是什么颜色，因为红色球在第0列白色球在第1列，所以0代表红色，1代表白色
        x = self.get_data_with_distribute(self.B[z])
        # 根据观测概率矩阵随机从第一个盒子选中一个颜色，x是颜色下标
        result = [x]
        for _ in range(T-1):        # 依次生成余下的状态和观测数据
            z = self.get_data_with_distribute(self.A[z]) # 根据转移概率，选择下一个状态，即下一次从哪个盒子种取球 #假设这里是2号盒子
            x = self.get_data_with_distribute(self.B[z]) # 选择取出什么颜色的球 假设x是白色，x就是1
            result.append(x)
        return result
    

    def forward_probability(self, X):
        '''
        根据给定的参数计算条件概率
        X: 观测数据 #球的颜色序列 红白红白白
        '''
        alpha = self.pi * self.B[:,X[0]]
        for x in X[1:]:
            alpha = np.matmul(alpha, A) * B[:, x]
        return alpha.sum() # alpha中存的是，t时刻，从四个盒子中取ot颜色的球，o1-ot-1颜色的概率值

    

if __name__ == "__main__":
    pi = np.array([.25, .25, .25, .25])
    A = np.array([
        [0,  1,  0, 0],
        [.4, 0, .6, 0],
        [0, .4, 0, .6],
        [0, 0, .5, .5]])
    B = np.array([
        [.5, .5],
        [.3, .7],
        [.6, .4],
        [.8, .2]])
    assert len(A) == len(pi)
    assert len(A) == len(B)
    hmm = HMM(B.shape[0], B.shape[1], pi, A, B)
    # print(hmm.generate(5))  # 生成5个数据
    seq = hmm.generate(5)
    print(seq)
    print(hmm.forward_probability(seq))